import { ModelInfo, ChunkMessage } from "./types";

export function parseModelName(modelInfo: ModelInfo) {
	const [ manufacturer, model ] = modelInfo.modelName.split('/');
	modelInfo.model = model;
	modelInfo.manufacturer = manufacturer;
	return modelInfo;
}

export function wrapMessages(messages: ChunkMessage[]): ChunkMessage[] {
	const chunks = [];
	const createDefaultMsg = () => {
		return {
			chunks: [],
			type: '' as 'text' |'reasoning' | 'usage' | 'error' | 'tool_calls',
			content: '',
			reasoning: '',
		} as ChunkMessage;
	};
	let currentMsg = createDefaultMsg();
	for(const msg of messages) {
		if(msg.type !== currentMsg.type) {
			if(currentMsg.type) {
				chunks.push(_mergeMessageFromChunks(currentMsg.chunks!));
			}
			currentMsg = createDefaultMsg();
			currentMsg.type = msg.type;
			currentMsg.chunks?.push(msg);
		} else {
			currentMsg.chunks?.push(msg);
		}
	}

	if(currentMsg.type && currentMsg.chunks && currentMsg.chunks.length > 0) {
		chunks.push(_mergeMessageFromChunks(currentMsg.chunks));
	}
	return chunks;
}

function _mergeMessageFromChunks(chunks: ChunkMessage[]): ChunkMessage {
	if(chunks.length === 1) {
		return chunks[0];	
	}
	const newChunk = {
		type: chunks[0].type,
		content: '',
		reasoning: '',
		tool_calls: [],
		finished: false
	} as ChunkMessage;

	for(const chunk of chunks) {
		if(chunk.type === 'text') {
			newChunk.content += chunk.content!;	
		}	else if	(chunk.type === 'reasoning') {
			newChunk.reasoning += chunk.reasoning!;
		}
		else if(chunk.type === 'tool_calls') {
			newChunk.tool_calls?.push(...chunk.tool_calls!);
		}
	}

	if(newChunk.type === 'tool_calls') {
		const toolCalls = newChunk.tool_calls ?? [];
		const newToolCalls = [];
		if(toolCalls.length > 0) {
			let firstToolCall = toolCalls[0];
			for(const toolCall of toolCalls) {
				toolCall.id = toolCall.id ?? firstToolCall.id;
				if(toolCall.id != firstToolCall.id) {
					newToolCalls.push(firstToolCall);
					firstToolCall = toolCall;
					continue;
				}
				firstToolCall.function.arguments += toolCall.function.arguments ?? '';
			}
			newToolCalls.push(firstToolCall);
		}
		newChunk.tool_calls = newToolCalls.map((item) => {
			try {
				item.function.args = JSON.parse(item.function.arguments ?? '{}');
			} catch (error) {
				console.log('parse tool call arguments failed', error);
			}
			return item;
		});
	}

	return newChunk;
}